import os
from functools import partial
from argparse import ArgumentParser
import random

import numpy as np

from tools.dataset_helper import open_img_np, plot_cmd
from tools.select_patch import crn_cmd
from tools.json_helper import json_helper


def sen2wkt(path: str) -> str:
    """
    将 sen12 的路径名提取为多模态唯一标识符
    """
    dir_, filename = os.path.split(path)
    dir_, region = os.path.split(dir_)
    _, roi = os.path.split(dir_)
    patch = filename.split(region+"_")[1]
    patch = os.path.splitext(patch)[0]
    region = region.split("_")[1]
    
    return str((roi, region, patch))


def wkt2sen(wkt: str, modal:str) -> str:
    
    if modal == "opt":
        mn = "s2"
    elif modal == "sar":
        mn = "s1"

    wkt = wkt.strip("()")
    wkt = wkt.split(",")
    for i in range(3):
        wkt[i] = wkt[i].strip("\' ")
    
    roi = wkt[0]
    region = mn + "_" + wkt[1]
    patch = roi.replace("rois", "ROIs")
    patch += "_" + region
    patch += "_" + wkt[2]

    return os.path.join(roi, region, patch+".png")


def open_sen12(
    src_file, *, scale=0, mean=False, dims=True, **kwargs
):
    """
    open sentinel1-2 image as numpy ndarray

    Param
    -----
    scale : change to range of [0 scale] after normalize
            no change range if scale == 0
    mean :  return mean of channels
    dims :  keep dims == 3

    Note
    -----
    plt.imread read png as float with [0, 1]
    otherwise int 
    """
    if scale == 0:
        max_ = np.array([0., 0., 0.])
    else:
        max_ = np.array([1., 1., 1.])

    img_np = open_img_np(
        src_file, max_=max_, scale=scale, mean=mean, dims=dims, **kwargs
    )

    return img_np


def get_sen12_pairs(filename : str, *, dset_dir='') -> tuple:
    '''
    根据某一模态的图片名获得两种模态的图片名(s1, s2)，如果dset_dir不为''，则返回的图片路径为相对于该值的相对路径
    '''
    if (dset_dir != '' and dset_dir in filename):
        filename = os.path.relpath(filename, dset_dir)
    new_filename = None
    ret = None
    if '_s1_' in filename:
        new_filename = filename.replace('_s1_', '_s2_')
        new_filename = new_filename.replace(
            os.sep+'s1', os.sep+'s2')
        ret = (filename, new_filename)
    elif '_s2_' in filename:
        new_filename = filename.replace('_s2_', '_s1_')
        new_filename = new_filename.replace(
            os.sep+'s2', os.sep+'s1')
        ret = (new_filename, filename)

    return ret


def get_img_union_from_img2dict(*args, ret_modal=1) -> list:
    '''
    输入由若干文件名列表组成地列表，输出他们在相同模态下并集的列表

    Params:
    -----
    ret_modal: 1、返回的文件名是s2的opt，0、返回的文件名是s1的sar
    '''
    ret = []
    for file_list in args:
        for filename in file_list:
            modal_filename = get_sen12_pairs(filename)[ret_modal]
            if modal_filename not in ret:
                ret.append(modal_filename)

    return ret


def get_sen12_subset(
    fold_list: list, *, rate: float=None, num: int=None
) -> dict:
    '''
    划分指定占比或者数量的子数据集

    Param:
    -----
    fold_list:表模态名的文件夹路径列表
    rate:子数据集占总数据集比例
    num:子数据集的数量，和rate参数仅一个有效

    Return:
    -----
    返回子数据集列表{'SAR': [], 'OPT': []}
    '''

    if rate != None:
        num = round(rate * len(fold_list))
    elif num == None:
        raise ValueError('both rate and num are None!')

    shuffle_ids = list(range(len(fold_list)))
    random.shuffle(shuffle_ids)
    subset_ids = shuffle_ids[:num]

    subset = {'SAR': [], 'OPT': []}
    for ids in subset_ids:
        pair = get_sen12_pairs(fold_list[ids])
        subset['SAR'].append(pair[0])
        subset['OPT'].append(pair[1])

    return subset


def get_subset_main (json_path, subset_path,*, num=None, rate=None):
    """
    json_path       : "存放场景图像文件列表文件的路径"
    subset_path     : "子数据集列表文件存储路径"
    
    """

    jh = json_helper(json_dir=json_path)
    dict_ = jh.get_dict()
    fold_list = get_img_union_from_img2dict(dict_['OPT'], dict_['SAR'])

    subset_dict = get_sen12_subset(fold_list, num=num, rate=rate) 
    jh_subset = json_helper(dict_=subset_dict)
    jh_subset.to_json(save_dir=subset_path)



# work as cmd 
# "crn" : get crn 
if __name__ == "__main__":

    p = input("select program --> ")

    if p == "crn":

        img2pt_file_list = crn_cmd(
            read_func=partial(open_sen12, scale=255, mean=True), 
            fn2wkt=sen2wkt, wkt2fn=wkt2sen, pt_range=[82, 174, 82, 174]
        )

        # python -m tools.sentinel12_helper -d "E:\datasets\sen1-2" -s "" -l ""  

        # python -m tools.sentinel12_helper -d "E:\datasets\sen1-2" -s "E:\workspace\SOMatch\json\sentinel1-2_harris\summer" -l "E:\workspace\SOMatch\json\sen12_list\rois1868_summer.json"  
        
        # python -m tools.sentinel12_helper -d "E:\datasets\sen1-2" -s "E:\workspace\SOMatch\json\sentinel1-2_harris\fall_200" -l "E:\workspace\SOMatch\json\sentinel1-2_harris\fall_200\fall_200.json"  
        
        # python -m tools.sentinel12_helper -d "E:\datasets\sen1-2" -s "E:\workspace\SOMatch\json\sentinel1-2_harris\spring_200" -l "E:\workspace\SOMatch\json\sentinel1-2_harris\spring_200\spring_200.json"   
        
        # python -m tools.sentinel12_helper -d "E:\datasets\sen1-2" -s "E:\workspace\SOMatch\json\sentinel1-2_harris\winter_200" -l "E:\workspace\SOMatch\json\sentinel1-2_harris\winter_200\winter_200.json"   
        
        # python -m tools.sentinel12_helper -d "E:\datasets\sen1-2" -s "E:\workspace\SOMatch\json\sentinel1-2_harris\summer_200" -l "E:\workspace\SOMatch\json\sentinel1-2_harris\summer_200\summer_200.json" 

        # python -m tools.sentinel12_helper -s "" -i "" "" -t   

        # python -m tools.sentinel12_helper -s "E:\workspace\SOMatch\json\sen12-test\fall_o2.8s1.1" -i "E:\workspace\SOMatch\json\sentinel1-2_harris\fall\OPT_HARRIS0.05.json" "E:\workspace\SOMatch\json\sentinel1-2_harris\fall\SAR_HARRIS0.05.json" -t 2.8 1.1

        # opt 1 sar 1.9 for fall
        # 

    if p == "plot":

        plot_cmd(
            wkt2fn=wkt2sen, 
            read_func=partial(open_sen12, scale=0, mean=True)
        )

        # python -m tools.sentinel12_helper -d "E:\datasets\sen1-2" -o "" -s "" -v "D:\datasets\client-data\plot\"

        # python -m tools.sentinel12_helper -d "E:\datasets\sen1-2" -o "E:\workspace\SOMatch\json\sentinel1-2_harris\summer\OPT_HARRIS0.05.json" -s "E:\workspace\SOMatch\json\sentinel1-2_harris\summer\SAR_HARRIS0.05.json" -v "D:\datasets\client-data\plot\sen_summer"

        # python -m tools.sentinel12_helper -d "E:\datasets\sen1-2" -o "E:\workspace\SOMatch\json\sentinel1-2_harris\fall_200\OPT_HARRIS0.05.json" -s "E:\workspace\SOMatch\json\sentinel1-2_harris\fall_200\SAR_HARRIS0.05.json" -v "D:\datasets\client-data\plot\sen_fall_200"

        # python -m tools.sentinel12_helper -d "E:\datasets\sen1-2" -o "E:\workspace\SOMatch\json\sentinel1-2_harris\spring_200\OPT_HARRIS0.05.json" -s "E:\workspace\SOMatch\json\sentinel1-2_harris\spring_200\SAR_HARRIS0.05.json" -v "D:\datasets\client-data\plot\sen_spring_200"

        # python -m tools.sentinel12_helper -d "E:\datasets\sen1-2" -o "E:\workspace\SOMatch\json\sentinel1-2_harris\winter_200\OPT_HARRIS0.05.json" -s "E:\workspace\SOMatch\json\sentinel1-2_harris\winter_200\SAR_HARRIS0.05.json" -v "D:\datasets\client-data\plot\sen_winter_200"

        # python -m tools.sentinel12_helper -d "E:\datasets\sen1-2" -o "E:\workspace\SOMatch\json\sentinel1-2_harris\summer_200\OPT_HARRIS0.05.json" -s "E:\workspace\SOMatch\json\sentinel1-2_harris\summer_200\SAR_HARRIS0.05.json" -v "D:\datasets\client-data\plot\sen_summer_200"


    if p == "subset":

        parser = ArgumentParser()
        parser.add_argument(
            '-j', '--json_path', help='json path of modal file lists', default=""
        )
        parser.add_argument(
            '-s', '--subset_path', help='folder to save results json', default=""
        )
        parser.add_argument(
            '-n', '--num', help='num of subset', type=int, default=0
        )
        args = parser.parse_args()

        get_subset_main(args.json_path, args.subset_path, num=args.num)

        # python -m tools.sentinel12_helper -s "" -j "" -n 200

        # python -m tools.sentinel12_helper -s "E:\workspace\SOMatch\json\sentinel1-2_harris\fall_200\fall_200.json" -j "E:\workspace\SOMatch\json\sen12_list\ROIs1970_fall.json" -n 200

        # python -m tools.sentinel12_helper -s "E:\workspace\SOMatch\json\sentinel1-2_harris\winter_200\winter_200.json" -j "E:\workspace\SOMatch\json\sen12_list\rois2017_winter.json" -n 200

        # python -m tools.sentinel12_helper -s "E:\workspace\SOMatch\json\sentinel1-2_harris\spring_200\spring_200.json" -j "E:\workspace\SOMatch\json\sen12_list\ROIs1158_spring.json" -n 200

        # python -m tools.sentinel12_helper -s "E:\workspace\SOMatch\json\sentinel1-2_harris\summer_200\summer_200.json" -j "E:\workspace\SOMatch\json\sen12_list\rois1868_summer.json" -n 200










